SAM_URL="http://172.16.159.9:25210/predict"

import requests
import numpy as np
import matplotlib.pyplot as plt
import cv2
from typing import List, Optional
import json
import os
import base64

# 设置matplotlib为非交互式后端
import matplotlib
matplotlib.use('Agg')  # 使用非交互式后端

def show_mask(mask, ax, random_color=False):
    if random_color:
        color = np.concatenate([np.random.random(3), np.array([0.6])], axis=0)
    else:
        color = np.array([30/255, 144/255, 255/255, 0.6])
    h, w = mask.shape[-2:]
    mask_image = mask.reshape(h, w, 1) * color.reshape(1, 1, -1)
    ax.imshow(mask_image)
    
def show_points(coords, labels, ax, marker_size=375):
    pos_points = coords[labels==1]
    neg_points = coords[labels==0]
    ax.scatter(pos_points[:, 0], pos_points[:, 1], color='green', marker='*', s=marker_size, edgecolor='white', linewidth=1.25)
    ax.scatter(neg_points[:, 0], neg_points[:, 1], color='red', marker='*', s=marker_size, edgecolor='white', linewidth=1.25)   
    
def show_box(box, ax):
    x0, y0 = box[0], box[1]
    w, h = box[2] - box[0], box[3] - box[1]
    ax.add_patch(plt.Rectangle((x0, y0), w, h, edgecolor='green', facecolor=(0,0,0,0), lw=2))

def image_to_base64(image_path: str) -> str:
    """将图像文件转换为base64字符串"""
    with open(image_path, "rb") as image_file:
        return base64.b64encode(image_file.read()).decode('utf-8')

def predict(
    image_path: str,
    point_coords: Optional[List[List[float]]] = None,
    point_labels: Optional[List[int]] = None,
    box: Optional[List[float]] = None,
    multimask_output: bool = True,
    mask_input: Optional[np.ndarray] = None
):
    # 确保文件存在
    if not os.path.exists(image_path):
        raise FileNotFoundError(f"Image file not found: {image_path}")
    
    # 创建请求体字典
    payload = {
        "image": image_to_base64(image_path),
        "multimask_output": multimask_output,
    }
    
    # 添加可选参数
    if point_coords is not None and len(point_coords) > 0:
        payload["point_coords"] = point_coords
    
    if point_labels is not None and len(point_labels) > 0:
        payload["point_labels"] = point_labels
    
    if box is not None and len(box) > 0:
        payload["box"] = box
    
    if mask_input is not None:
        payload["mask_input"] = mask_input.tolist()
    
    print("Sending payload with image size:", len(payload["image"]), "chars")
    
    # 发送请求
    try:
        response = requests.post(
            SAM_URL,
            json=payload
        )
        
        # 检查响应状态
        if response.status_code != 200:
            error_msg = f"Request failed with status {response.status_code}"
            if response.text:
                error_msg += f": {response.text}"
            raise Exception(error_msg)
        
        # 处理响应
        result = response.json()
        masks = np.array(result["masks"])
        scores = np.array(result["scores"])
        logits = np.array(result["logits"]) if result.get("logits") else None
        
        return masks, scores, logits
    
    except requests.exceptions.RequestException as e:
        raise Exception(f"Network error: {str(e)}")
    except json.JSONDecodeError:
        raise Exception("Failed to parse server response")
    except KeyError as e:
        raise Exception(f"Missing expected field in response: {str(e)}")

def save_mask_result(image, mask, points, labels, score, filename):
    """保存掩码结果到文件"""
    plt.figure(figsize=(10,10))
    plt.imshow(image)
    show_mask(mask, plt.gca())
    show_points(points, labels, plt.gca())
    plt.title(f"Mask, Score: {score:.3f}", fontsize=18)
    plt.axis('off')
    plt.savefig(filename, bbox_inches='tight')
    plt.close()
    print(f"Saved result to {filename}")

def save_box_result(image, mask, box, filename):
    """保存框选结果到文件"""
    plt.figure(figsize=(10,10))
    plt.imshow(image)
    show_mask(mask, plt.gca())
    show_box(box, plt.gca())
    plt.axis('off')
    plt.savefig(filename, bbox_inches='tight')
    plt.close()
    print(f"Saved result to {filename}")

# 示例用法
if __name__ == "__main__":
    try:
        # 加载图像
        image_path = "images/female_skirt.jpg"
        print("Image path:", os.path.abspath(image_path))
        
        # 确保图像存在
        if not os.path.exists(image_path):
            print(f"Error: Image file not found at {os.path.abspath(image_path)}")
            exit(1)
        
        # 读取图像用于可视化
        image = cv2.imread(image_path)
        if image is None:
            print(f"Failed to load image: {image_path}")
            exit(1)
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        
        # 示例1: 点选对象
        print("Testing point selection...")
        masks, scores, logits = predict(
            image_path,
            point_coords=[[500, 375]],  # 点坐标
            point_labels=[1]             # 点标签
        )
        
        # 保存结果而不是显示
        for i, (mask, score) in enumerate(zip(masks, scores)):
            save_mask_result(
                image, 
                mask, 
                np.array([[500, 375]]), 
                np.array([1]), 
                score, 
                f"point_mask_{i+1}.png"
            )
        
        # 示例2: 框选对象
        print("Testing box selection...")
        masks, _, _ = predict(
            image_path,
            box=[425, 600, 700, 875],  # 边界框 [x1, y1, x2, y2]
            multimask_output=False
        )
        
        save_box_result(
            image,
            masks[0],
            [425, 600, 700, 875],
            "box_mask.png"
        )
    
    except Exception as e:
        print(f"Error: {str(e)}")
        # 打印详细错误信息
        import traceback
        traceback.print_exc()